import errno
import logging
import select
import socket
from typing import Dict
from typing import Optional
from typing import Sequence
from typing import Tuple

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s][%(processName)s:%(process)d:%(threadName)s][%(filename)s:%(lineno)d] - %(message)s",
)
logger = logging.getLogger(__name__)
SHOULDCONTINUE = (errno.EWOULDBLOCK, errno.EAGAIN)
DISCONNECTED = (
    errno.ECONNRESET,
    errno.ENOTCONN,
    errno.ESHUTDOWN,
    errno.ECONNABORTED,
    errno.EPIPE,
    errno.EBADF,
)


class Socket(object):
    """
    对原始socket的封装
    """

    def __init__(
        self, address_or_socket: str | socket.socket, peer_addr=None, listen=False
    ):
        """
        如果listen=True, 则传入服务端地址(str类型);
        否则, 代表一个建立连接的socket
        """
        self.listen = listen
        if self.listen:
            assert isinstance(address_or_socket, str)
            self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
            self._sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 0)
            self._sock.setblocking(False)
            self.address = address_or_socket
        else:
            assert isinstance(address_or_socket, socket.socket)
            self._sock = address_or_socket

        self.fd = self._sock.fileno()
        if self.listen:
            self._listen(self.address)
        if peer_addr is None:
            peer_addr = ("unknown", -1)
        self.peer_addr = peer_addr

    def _listen(self, addr):
        if not self.listen:
            logger.warning("can not listen to socket")
            return
        host, port = addr.split(":")
        self._sock.bind((host, int(port)))
        logger.info(f"listening on {addr} ...")
        backlog = 8
        self._sock.listen(backlog)

    def accept(self) -> Optional["Socket"]:
        if not self.listen:
            raise Exception("can not accept")
        try:
            conn_sock, peer_addr = self._sock.accept()
            return Socket(conn_sock, peer_addr=peer_addr)
        except socket.error as e:
            if e.args[0] in SHOULDCONTINUE:
                return None
            elif e.args[0] == errno.ECONNABORTED:
                # errno.ECONNABORTED: 客户端发送了rst
                return None
            logger.error("accept error")
            raise

    def close(self):
        try:
            self._sock.shutdown(socket.SHUT_RDWR)
        except socket.error as e:
            if e.args[0] in (errno.ENOTCONN, errno.EBADF):
                # errno.ENOTCONN: 关闭一个已经被关闭了的连接
                # errno.EBADF: 描述符失效
                return
            logger.error(f"socket close error: {e}")
            raise

    def read(self) -> Tuple[bytes, bool]:
        """
        非阻塞读取数据
        返回(data, is_closed), is_closed表示是否已经读完(连接关闭时, 是一个可读事件, 读取的数据为空字节序列)
        """
        is_closed = False
        try:
            data = self._sock.recv(1024)
            if not data:
                is_closed = True
            return data, is_closed
        except socket.error as e:
            logger.debug(f"socket recv error: {e}")
            if e.args[0] in SHOULDCONTINUE:
                return b"", True
            elif e.args[0] in DISCONNECTED:
                return b"", True
            logger.error("socket recv error")
            raise

    def write(self, data: bytes) -> Tuple[int, bool]:
        """
        非阻塞方式写入数据
        返回 (写入的长度, 连接是否关闭)
        """
        try:
            write_length = self._sock.send(data)
            return write_length, False
        except socket.error as e:
            if e.args[0] in SHOULDCONTINUE:
                return 0, False
            elif e.args[0] in DISCONNECTED:
                return 0, True
            logger.error("socket send error")
            raise


class Channel(object):
    """
    连接通道
    对文件描述符的抽象.
    """

    def __init__(self, loop: "EventLoop", fd: int):
        self._loop = loop
        self.fd = fd

        self.need_write = False
        self.need_read = False
        self.writable = False
        self.readable = False
        self.error = False

        self.read_callback = None
        self.write_callback = None
        self.error_callback = None

    def register(self):
        """
        注册到事件循环中
        """
        self._loop.add(self)

    def disable(self):
        """
        关闭对该channel的监测
        """
        self.need_write = False
        self.need_read = False

    def close(self):
        self._loop.remove(self)

        self.read_callback = None
        self.write_callback = None
        self.error_callback = None

    def set_read_callback(self, callback):
        self.read_callback = callback

    def set_write_callback(self, callback):
        self.write_callback = callback

    def set_error_callback(self, callback):
        self.error_callback = callback

    def handle_event(self):
        """
        channel就绪后执行的回调
        """
        logger.debug("handle_event")
        if self.readable and self.read_callback:
            self.read_callback()

        if self.writable and self.write_callback:
            self.write_callback()

        if self.error and self.error_callback:
            self.error_callback()


class Poller(object):
    def __init__(self):
        # fd到channel的映射
        self._fd_channel: Dict[int, Channel] = {}

    def update_channel(self, channel: Channel):
        self._fd_channel[channel.fd] = channel

    def remove_channel(self, channel: Channel):
        if channel.fd in self._fd_channel:
            del self._fd_channel[channel.fd]

    def poll(self) -> Sequence[Channel]:
        """
        拉取所有就绪的channel
        """
        ready_list = []

        r_list = []
        w_list = []
        x_list = []

        if not self._fd_channel:
            return []

        for fd in self._fd_channel:
            channel = self._fd_channel[fd]
            if channel.need_read:
                r_list.append(fd)

            if channel.need_write:
                w_list.append(fd)

            if channel.need_read or channel.need_write:
                # ?
                x_list.append(fd)

        if not r_list and not w_list and not x_list:
            logger.debug("no event")
            return []

        try:
            r_list, w_list, x_list = select.select(r_list, w_list, x_list, None)
            logger.debug("select end")
        except select.error as e:
            if e.args[0] == errno.EINTR:
                return []
            logger.info(f"exception: {e}")
            raise

        for rfd in r_list:
            channel = self._fd_channel[rfd]
            channel.readable = True
            ready_list.append(channel)

        for wfd in w_list:
            channel = self._fd_channel[wfd]
            channel.writable = True
            ready_list.append(channel)

        for efd in x_list:
            channel = self._fd_channel[efd]
            channel.error = True
            ready_list.append(channel)

        return ready_list


class EventLoop(object):
    def __init__(self):
        self._poller = Poller()

    def add(self, channel):
        self._poller.update_channel(channel)

    def remove(self, channel):
        self._poller.remove_channel(channel)

    def run_loop(self):
        while True:
            try:
                ready_list = self._poller.poll()
                for channel in ready_list:
                    channel.handle_event()
            except KeyboardInterrupt:
                break


class Acceptor(object):
    """
    管理tcp服务端监听socket
    """

    def __init__(self, loop, addr):
        self._evloop = loop
        # 创建socket
        self.listen_sock = Socket(addr, listen=True)

        # 创建通道
        self.listen_channel = Channel(self._evloop, self.listen_sock.fd)
        self.listen_channel.need_read = True
        self.listen_channel.set_read_callback(self.handle_read)
        self.listen_channel.register()

    def set_new_connection_callback(self, callback):
        """
        设置新连接回调
        回调函数签名: callback(conn_sock: Socket) -> None
        """
        self.set_new_connection_callback = callback

    # 可读(有新的连接时)回调
    def handle_read(self):
        conn_sock = self.listen_sock.accept()
        if conn_sock is None:
            return

        logger.debug(f"new connection: {conn_sock.peer_addr}")
        if self.set_new_connection_callback:
            self.set_new_connection_callback(conn_sock)


class TcpConnection(object):
    """
    服务端, tcp连接
    对socket.accept()得到的conn_sock的抽象
    """

    def __init__(self, loop, conn_sock: Socket):
        self._evloop = loop
        self.conn_sock = conn_sock
        self.channel = Channel(self._evloop, self.conn_sock.fd)
        self.channel.set_read_callback(self.handle_read)
        self.channel.set_write_callback(self.handle_write)
        self.channel.set_error_callback(self.handle_error)
        self.channel.need_read = True
        # 事件循环监听这一个新连接的事件
        self.channel.register()

        # tcp连接唯一标识符
        self.id = f"{self.conn_sock.peer_addr[0]}:{self.conn_sock.peer_addr[1]}|{self.conn_sock.fd}"

        # 回调函数签名: callback(conn_id: str) -> None
        self.close_callback = None
        # 回调函数签名: callback() -> None
        self.write_complete_callback = None

        self.read_buffer = b""
        self.write_buffer = b""

        # handler.handle_read 回调函数签名: callback(conn: TcpConnection, readbuf: bytes) -> None
        self.handler = None

    # 内部回调(所有应用都一样的处理)

    def handle_read(self):
        logger.debug("TcpConnection handle_read")
        data, is_end = self.conn_sock.read()
        if is_end:
            self.handle_close()
            return

        self.read_buffer += data
        logger.debug(f"read buffer: {self.read_buffer}")

        # 临时这么处理, 为了测试
        if self.handler:
            self.handler.handle_read(self, self.read_buffer)
            self.read_buffer = b""

    def handle_write(self):
        logger.debug("TcpConnection handle_write")
        write_length, is_closed = self.conn_sock.write(self.write_buffer)
        if is_closed:
            self.handle_close()
            return

        if write_length == len(self.write_buffer):
            # 一次写完了
            self.channel.need_write = False
            self.write_buffer = b""

            if self.write_complete_callback:
                self.write_complete_callback()
        else:
            # 部分写入
            self.write_buffer = self.write_buffer[write_length:]

    def handle_close(self):
        logger.debug("TcpConnection handle_close")
        self.read_buffer = b""
        self.channel.disable()
        self.conn_sock.close()

        if self.close_callback:
            self.close_callback(self.id)

    def handle_error(self):
        logger.debug("TcpConnection handle_error")

    # 需要外部设置的回调(根据应用需要设置)
    def set_close_callback(self, callback):
        self.close_callback = callback

    def set_write_complete_callback(self, callback):
        self.write_complete_callback = callback

    def send(self, data: bytes):
        write_length, is_closed = self.conn_sock.write(data)
        if is_closed:
            self.handle_close()
            return

        if write_length < len(data):
            # 只写入了部分
            self.write_buffer += data[write_length:]
            self.channel.need_write = True
        elif write_length == len(data) and self.write_complete_callback:
            # 一次发送完, 就不用设置need_write标志(也就不会经过handle_write回调函数)
            self.write_complete_callback()


class TcpServer(object):
    """
    tcp服务端
    """

    def __init__(self, addr):
        self._evloop = EventLoop()
        self._acceptor = Acceptor(self._evloop, addr)
        self._acceptor.set_new_connection_callback(self.handle_new_connection)

        self._id_conn: Dict[str, TcpConnection] = {}

    def handle_new_connection(self, conn_sock: Socket):
        logger.debug("handling new connection ...")
        new_conn = TcpConnection(self._evloop, conn_sock)
        new_conn.set_close_callback(self.handle_remove_connection)
        new_conn.set_write_complete_callback(self.handle_write_complete)
        self._id_conn[new_conn.id] = new_conn

        if self.handler_class:
            new_conn.handler = self.handler_class()

    def handle_remove_connection(self, conn_id: str):
        logger.debug("handling remove connection ...")
        if conn_id in self._id_conn:
            del self._id_conn[conn_id]
            logger.debug(f"tcp connection: {conn_id} deleted.")

    def handle_write_complete(self):
        logger.debug("handling write_complete ...")

    def add_handler(self, handler_class):
        self.handler_class = handler_class

    def run(self):
        # 将channel添加到事件循环
        self._evloop.run_loop()


class EchoHandler(object):
    def handle_read(self, conn: TcpConnection, read_buffer: bytes):
        conn.send(read_buffer)


if __name__ == "__main__":
    server = TcpServer("localhost:8080")
    server.add_handler(EchoHandler)
    server.run()
